
import zipfile
import os
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

# 定义下载文件的函数
def download_file(url, local_filename):
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()  # 检查请求是否成功
        with open(local_filename, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)
        print(f"Downloaded {local_filename}")
        return local_filename
    except requests.RequestException as e:
        print(f"Failed to download file {url}: {e}")
        return None

# 定义解压文件的函数
def extract_zip(file_path, extract_to):
    with zipfile.ZipFile(file_path, 'r') as zip_ref:
        zip_ref.extractall(extract_to)
        print(f"Extracted {file_path} to {extract_to}")

# 下载并解压所有CSV文件
base_url = "http://storage.googleapis.com/books/ngrams/books/googlebooks-chi-sim-all-3gram-20090715-"
# base_url = "http://storage.googleapis.com/books/ngrams/books/googlebooks-chi-sim-all-2gram-20090715-"
# base_url = "http://storage.googleapis.com/books/ngrams/books/googlebooks-chi-sim-all-1gram-20090715-"
# output_dir = "extracted_files1"
# output_dir = "extracted_files2"
output_dir = "extracted_files3"
os.makedirs(output_dir, exist_ok=True)

# 使用线程池并行下载文件
urls = [f"{base_url}{i}.csv.zip" for i in range(100)]
# urls = [f"{base_url}{i}.csv.zip" for i in range(100)]
# urls = [f"{base_url}{i}.csv.zip" for i in range(10)]
with ThreadPoolExecutor(max_workers=10) as executor:
    future_to_url = {executor.submit(download_file, url, os.path.join(output_dir, os.path.basename(url))): url for url in urls}
    
    for future in as_completed(future_to_url):
        url = future_to_url[future]
        try:
            data = future.result()
            if data:
                extract_zip(data, output_dir)
                os.remove(data)  # 删除已解压的ZIP文件
        except Exception as exc:
            print(f'{url} generated an exception: {exc}')
